import triton
import triton.language as tl
import torch
import torch_npu

from torch._inductor.runtime.triton_helpers import math as tl_math
from torch._inductor.runtime import triton_helpers
from .util import triton_helper_compile, compiled_kernel_list

@triton.jit
def fused_kernel4(in_ptr0, in_ptr1, in_ptr2, out_ptr3, xnumel, rnumel):
    xnumel = 288
    XBLOCK: tl.constexpr = 1
    rnumel = 9
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = tl.full([1], xoffset, tl.int32)
    xmask = tl.full([RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[:]
    roffset = 0
    rmask = rindex < rnumel
    r2 = rindex
    x3 = xindex
    x0 = xindex % 9
    tmp0 = tl.load(in_ptr0 + (r2 + (9*x3)), rmask, other=0.0)
    tmp7 = tl.load(in_ptr1 + (x0), None, eviction_policy='evict_last')
    tmp14 = tl.load(in_ptr2 + (r2), rmask, eviction_policy='evict_last', other=0.0)
    tmp1 = r2 + ((-1)*x0)
    tmp2 = tl.full([1], 1, tl.int32)
    tmp3 = tmp1 >= tmp2
    tmp4 = -3.4028234663852886e+38
    tmp5 = 0.0
    tmp6 = tl.where(tmp3, tmp4, tmp5)
    tmp8 = tmp7.to(tl.float32)
    tmp9 = r2
    tmp10 = tmp9.to(tl.float32)
    tmp11 = tmp10 > tmp8
    tmp12 = tmp11.to(tl.float32)
    tmp13 = tmp6 * tmp12
    tmp15 = tmp14.to(tl.float32)
    tmp16 = tmp13 + tmp15
    tmp17 = tmp16 == tmp5
    tmp18 = tl.where(tmp17, tmp4, tmp13)
    tmp19 = tmp0 + tmp18
    tmp20 = float("-inf")
    tmp21 = tmp19 == tmp20
    tmp22 = tmp21 == 0
    tmp23 = tmp22.to(tl.int32)
    tmp24 = (tmp23 != 0)
    tmp25 = tl.broadcast_to(tmp24, [RBLOCK])
    tmp27 = tl.where(rmask, tmp25, 0)
    tmp28 = triton_helpers.promote_to_tensor(triton_helpers.any(tmp27, 0))
    tmp29 = tl.broadcast_to(tmp19, [RBLOCK])
    tmp31 = tl.where(rmask, tmp29, float("-inf"))
    tmp32 = triton_helpers.promote_to_tensor(triton_helpers.max2(tmp31, 0))
    tmp33 = tmp19 - tmp32
    tmp34 = tl_math.exp(tmp33)
    tmp35 = tl.broadcast_to(tmp34, [RBLOCK])
    tmp37 = tl.where(rmask, tmp35, 0)
    tmp38 = triton_helpers.promote_to_tensor(tl.sum(tmp37, 0))
    tmp39 = tmp28 == 0
    tmp40 = tmp34 / tmp38
    tmp41 = tl.where(tmp39, tmp5, tmp40)
    tl.store(out_ptr3 + (r2 + (9*x3)), tmp41, rmask)

def test_compile_for_fused_kernel4():
    signature = {0: '*fp32', 1: '*i32', 2: '*i32', 3: '*fp32', 4: 'i32', 5: 'i32'}
    triton_helper_compile(func=fused_kernel4, signature=signature)

compiled_kernel_list["fused_kernel4"] = test_compile_for_fused_kernel4